import gradio as gr
from langchain_openai import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage
import time

api_key = '503687a4a1b3de564e8a65c9da87b3da.VOdqJVr0l2WX69Ie'

llm = ChatOpenAI(temperature=0.95, model_name='glm-4', openai_api_key=api_key,
                 openai_api_base='https://open.bigmodel.cn/api/paas/v4/', streaming=True)


def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)


def add_message(history, message):
    for x in message["files"]:
        history.append(((x,), None))
    if message["text"] is not None:
        history.append((message["text"], ''))
    return history, message, gr.MultimodalTextbox(value=None, interactive=False)


def bot(history):
    # [['你是谁？', "**That's cool!**"], ['我是谁？', "**That's cool!**"], ['你能做什么？', '']]
    history_langchain_format = []
    message = history[-1][0]
    for human, ai in history:
        if ai is not None and ai != '':
            history_langchain_format.append(HumanMessage(content=human))
            history_langchain_format.append(AIMessage(content=ai))

    history_langchain_format.append(HumanMessage(content=message))
    # 返回类型是一个langchain_core.messages.ai.AIMessage
    glm_response = llm(history_langchain_format)
    history[-1][1] = ""
    for character in glm_response.content:
        history[-1][1] += character
        time.sleep(0.05)
        yield history


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            top_p = gr.Slider(0.00, 1.00, label="top_p", interactive=True)
            temperature = gr.Slider(0.00, 1.00, label="temperature", interactive=True)
            top_k = gr.Slider(1, 20, label="top_k", interactive=True)
            out_input_length = gr.Slider(1, 2048, label="out_input_length", interactive=True)
            image = gr.Image(label="Choose an image...")
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                [],
                elem_id="chatbot",
                bubble_full_width=False
            )

            with gr.Row():
                btn1 = gr.Button("Retry")
                btn2 = gr.Button("Undo")
                btn3 = gr.Button("Clear")

            chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"],
                                              placeholder="Enter message or upload file...", show_label=False)
            chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
            bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
            bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

            chatbot.like(print_like_dislike, None, None)

demo.queue()
demo.launch(share=True)
